{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from tqdm.notebook import tqdm\n",
    "import sisl\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In TBtrans and TranSiesta one is capable of performing real-space transport calculations by using real-space self-energies (see [here](https://doi.org/10.1103/PhysRevB.100.195417)).  \n",
    "Currently, the real-space self-energy calculation *has* to be performed in `sisl` since it is not implemented in TranSiesta.\n",
    "\n",
    "A real-space self-energy is a $\\mathbf k$-averaged self-energy which can emulate *any* 2D or 3D electrode. I.e. for an STM junction a tip and a surface. In such a system, the surface could be modelled using the real-space self-energy to remove mirror effects of the STM tip periodic images. This is important since the distance between periodic images disturbs the calculation due to long range potential effects.\n",
    "\n",
    "The basic principle for calculating the real-space self-energy is the Brillouin zone integral:\n",
    "\\begin{equation}\n",
    "   \\mathbf G_{\\mathcal R}(E) = \\int_{\\mathrm{BZ}}\\mathbf G_\\mathbf k(E) = \\int_{\\mathrm{BZ}} \\big[E \\mathbf S_{\\mathbf k} - \\mathbf H_{\\mathbf k}\\big]^{-1}\n",
    "\\end{equation}\n",
    "Where $\\mathbf G_{\\mathcal R}$ is the real-space Green function, *E* is the energy, and $\\mathbf G_\\mathbf k$, $\\mathbf S_{\\mathbf k}$ and $\\mathbf H_{\\mathbf k}$ are the Green function, overlap matrix and Hamiltonian, respectively, at a given **k**-point. Integration is done over the entire Brillouin-zone ($\\mathrm{BZ}$). \n",
    "\n",
    "In this example, we will construct a STM tip probing a graphene flake.\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "This example is rather complicated and is the reason why basically everything is already done for you. Please try and understand each step.\n",
    "\n",
    "We start by creating the graphene tight-binding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphene = sisl.geom.graphene(orthogonal=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Graphene tight-binding parameters\n",
    "on, nn = 0, -2.7\n",
    "H_minimal = sisl.Hamiltonian(graphene)\n",
    "H_minimal.construct([[0.1, 1.43], [on, nn]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the minimal graphene unit-cell (here orthogonal) is created, we now turn to the calculation of the real-space self-energy.  \n",
    "The construction of this object is somewhat complicated and has a set of required input options:\n",
    "- `object`: the Hamiltonian\n",
    "- `semi_axis`: which axis to use for the recursive calculation of the self-energy\n",
    "- `k_axes`: which axes to integrate in the Brillouin zone\n",
    "- `unfold`: how many times the `object` needs to be unfolded along each lattice vector, this is an integer vector of length three"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# object = H_minimal\n",
    "# semi_axis = 0, x-axis uses recursive self-energy calculation\n",
    "# k_axes = 1, y-axis uses a Brillouin zone integral\n",
    "# unfold = (10, 10, 1), the full real-space Green function is equivalent to the system\n",
    "#                       H_minimal.tile(10, 0).tile(10, 1)\n",
    "RSSE = sisl.RealSpaceSE(H_minimal, 0, 1, (10, 10, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create the real-space self-energy.  \n",
    "In TBtrans (and TranSiesta) the electrode atomic indices *must* be in consecutive order.\n",
    "This is a little troublesome since the natural order in a device would be in order according to $x$, $y$ or $z$. To create the correct order we extract the *real-space coupling* matrix which is where the real-space self-energy would live. The self-energy ($\\boldsymbol\\Sigma^{\\mathcal R}$) is calculated using:\n",
    "\\begin{align}\n",
    "   \\mathbf G^{\\mathcal R}(E) &= \\big[E \\mathbf S^{\\mathcal R} - \\mathbf H^{\\mathcal R} - \\boldsymbol\\Sigma^{\\mathcal R}\\big]^{-1}\n",
    "   \\\\\n",
    "   \\boldsymbol\\Sigma^{\\mathcal R}(E) &= E \\mathbf S^{\\mathcal R} - \\mathbf H^{\\mathcal R} - \\Big[\\mathbf G^{\\mathcal R}(E)\\Big]^{-1}.\n",
    "\\end{align}\n",
    "Where $\\mathbf S^{\\mathcal R}$ and $\\mathbf H^{\\mathcal R}$ are the real-space overlap matrix and Hamiltonian, respectively. Another way to calculate the self-energy would be to transfer the Green function from the infinite bulk into the region of interest:\n",
    "\\begin{equation}\n",
    "   \\boldsymbol\\Sigma^{\\mathcal R} = \\mathbf V_{\\infty\\setminus\\mathcal R,\\mathcal R}\\mathbf G_{\\mathcal R,\\infty\\setminus\\mathcal R}\\mathbf V_{\\infty\\setminus\\mathcal R,\\mathcal R}.\n",
    "\\end{equation}\n",
    "Where $\\mathbf V_{\\infty\\setminus\\mathcal R,\\mathcal R}$ is the **to_be_filled** and $\\mathbf G_{\\mathcal R,\\infty\\setminus\\mathcal R}$ is the **to_be_filled**. From the second equation it is obvious that the self-energy only lives on the boundary that $\\mathbf V_{\\infty\\mathcal R,\\mathcal R}$ couples to. Exactly this region is extracted using `real_space_coupling` as below. Take some time to draw a simple 2D lattice coupling and confirm the area that the real-space self-energies couple to.\n",
    "*HINT:* Draw the unit-cell for a 4x4 square lattice with nearest-neighbour couplings, then mark the atoms that couple to atoms outside the unit-cell.  \n",
    "\n",
    "----\n",
    "\n",
    "In this example, we also retrieve the indices for the electrode atoms, those that connect *out to the infinite plane*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H_elec, elec_indices = RSSE.real_space_coupling(ret_indices=True)\n",
    "H_elec.write('GRAPHENE.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above yields the electrode region which contains the self-energies. The full device region is nothing but the `H_minimal` tiled $10\\times10$ times with an attached STM tip on top. Here we need to arrange the electrode atoms first, and then the final atoms of the device region after. The `real_space_parent` method returns the Hamiltonian that obeys the *unfolded* size - in this case $10\\times10$ times larger. One should always use this method to get the correct device order of atoms, since the order of tiling is determined by the `semi_axes` and `k_axis` arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H = RSSE.real_space_parent()\n",
    "# Create the true device by re-arranging the atoms\n",
    "indices = np.arange(len(H))\n",
    "indices = np.delete(indices, elec_indices)\n",
    "# first electrodes, then rest of device\n",
    "indices = np.concatenate([elec_indices, indices])\n",
    "# Now re-arange matrix\n",
    "H = H.sub(indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we need to add the STM tip. Here we simply add a gold atom and manually add the hoppings. Since this is tight-binding, we have full control over the self-energy and potential landscape. Therefore we don't need to extend the electrode region to screen-off the tip region. In DFT systems a screening region is required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "STM = sisl.Geometry([0, 0, 0], atoms=sisl.Atom('Au', R=1.0001), lattice=sisl.Lattice([10, 10, 1], nsc=[1, 1, 3]))\n",
    "H_STM = sisl.Hamiltonian(STM)\n",
    "H_STM.construct([(0.1, 1.1), (0, -0.75)])\n",
    "H_STM.write('STM.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mid_xyz = H.geometry.center()\n",
    "idx = H.close(mid_xyz, R=1.33)[0]\n",
    "H_device = H.add(H_STM, offset=H.geometry.xyz[idx] - H_STM.geometry.xyz[0] + [0, 0, 2])\n",
    "idx_STM = len(H_device) - 1\n",
    "idx = H_device.close(na, R=(0.1, 2.25))[1][0]\n",
    "H_device[idx_STM, idx] = -0.1\n",
    "H_device[idx, idx_STM] = -0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H_device.write('DEVICE.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we can run calculations, we need to create the real-space self-energy for the graphene flake in `sisl`.\n",
    "Since the algorithm is not implemented in TBtrans (nor TranSiesta) it needs to be done here.\n",
    "\n",
    "This is somewhat complicated since the files require a specific order. For ease, this tutorial implements it for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A real-space transport calculation ONLY needs the Gamma-point\n",
    "gamma = sisl.MonkhorstPack(H_elec, [1, 1, 1])\n",
    "# Energy-contour\n",
    "dE = 0.04\n",
    "E = np.arange(-2, 2 + dE / 2, dE)\n",
    "sisl.io.tableSile(\"contour.E\", 'w').write_data(E, np.zeros(E.size) + dE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now create the file (should take around 3-4 minutes)\n",
    "eta = 0.001 * 1j\n",
    "with sisl.io.tbtgfSileTBtrans(\"GRAPHENE.TBTGF\") as f:\n",
    "    f.write_header(gamma, E + eta)\n",
    "    for ispin, new_k, k, e in tqdm(f, unit=\"rsSE\"):\n",
    "        if new_k:\n",
    "            f.write_hamiltonian(H_elec.Hk(format='array', dtype=np.complex128))\n",
    "        SeHSE = RSSE.self_energy(e + eta, bulk=True, coupling=True)\n",
    "        f.write_self_energy(SeHSE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "- Calculate transport, density-of-states and bond-transmissions.  \n",
    "   Please search the manual on how to edit the `RUN.fdf` according to the following:\n",
    "   - Force `tbtrans` to use the generated `TBTGF` file. This is the same as in [TB_07](../TB_07/run.ipynb) example, i.e. using the `out-of-core` tag\n",
    "   - Force `tbtrans` to use an energy grid defined in an external file (`contour.E`). Search the `tbtrans` manual.\n",
    "   \n",
    "- Plot the bond-transmissions and check their symmetry, does the symmetry depend on the injection point?\n",
    "- Is there a particular reason for choosing `semi_axis` and `k_axes` as they are chosen? Or could they be swapped?\n",
    "\n",
    "- **TIME**: Redo the calculations using three electrodes (left/right/tip) using *k*-points. Converge transmission and then plot the bond-transmissions.  \n",
    "   Do they look as the real-space calculation? If they are the same, why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tbtrans RUNfix.fdf > TBT.out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tbt = sisl.get_sile('siesta.TBT.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learned lessons\n",
    "\n",
    "- Creation of real-space self-energies\n",
    "- Manual specification of energy points for TBtrans\n",
    "- Manual creation of self-energy files for TBtrans\n",
    "- Forcing TBtrans to use an existing **TBTGF** file on disk, previously created with `sisl`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
